from ultralytics import YOLO
import time
import cv2
import os
import numpy as np

start1 = time.time()

model = YOLO('last.pt')  # load a pretrained model (recommended for training)
model.info()
start2 = time.time()
# img = cv2.imread('F:/dataset_yimu/luoding_test/Image_20231018141420717.jpg')
# result = model.predict(source=img)
# print(len(result))
# for r in result:
#     print(r)
#     # cv2.imshow('box', r.orig_img)
#     # cv2.waitKey()

# file_path = 'F:/dataset_yimu/luoding_test/Image_20231018141420717.jpg' # luoding
file_path = 'F:/dataset_yimu/luoding_test/Image_20231018142138496.jpg' # wu luoding
# file_path = 'F:/dataset_yimu/hangtiankegong/test_save/0.jpg' # nothing
results = model(file_path)
print(len(results))
for r in results:
    print()
    # im_array = r.plot()
    # cv2.imshow('box', im_array)
    # cv2.waitKey()
end = time.time()
print(end - start1)
print(end - start2)

# model = YOLO('../runs/detect/train69/weights/last.pt')
# file_path = 'F:/dataset_yimu/luoding_test/Image_20231018141420717.jpg' # luoding
# file_path = 'F:/dataset_yimu/luoding_test/Image_20231018142138496.jpg' # wu luoding
# # file_path = 'F:/dataset_yimu/hangtiankegong/test_save/0.jpg' # nothing
# img = cv2.imread(file_path)
# def get_processed_img(img):
#     result = model.predict(source=img)
#     result = result[0]
#     print(result)
#     # processed_img = result.
#     # return processed_img, ld, detected_id
# get_processed_img(img)
# # img_r, _, _ = get_processed_img(img)
# cv2.imshow('box', img)
# cv2.waitKey()